package com.rhino.security.handler;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.rhino.security.constants.SecurityConstants;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.servlet.http.HttpServletRequest;

/**
 * Rest 接口返回到前端统一格式；
 * 返回的响应体增强处理器；
 */
@Slf4j
@RestControllerAdvice
public class JwtResponseBodyHanlder implements ResponseBodyAdvice<Object> {

    // 新生成的jwtToken
    private Object token;

    @Override
    public boolean supports(MethodParameter methodParameter, Class<? extends HttpMessageConverter<?>> aClass) {
        ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = servletRequestAttributes.getRequest();
        Object attribute = request.getAttribute(SecurityConstants.INITIATIVE_JWT_TOKEN);
        if (attribute != null) {
            token = attribute;
            return true;
        }
        return false;
    }

    @Override
    public Object beforeBodyWrite(Object object, MethodParameter methodParameter,
                                  MediaType mediaType, Class<? extends HttpMessageConverter<?>> aClass,
                                  ServerHttpRequest servletRequest, ServerHttpResponse serverHttpResponse) {


        if (object != null) {

            if(object instanceof String){
                return object;
            }
            //Controller 层返回的对象统一包装一下
            try {
                String body = JSON.toJSONString(object);
                JSONObject jsonObject = JSON.parseObject(body);
                jsonObject.put(SecurityConstants.JWT_TOKEN_VO, token);
                return jsonObject;
            } catch (Exception exception) {
                // 如果返回的不是json格式的值，那么则返回值不做处理。只捕获异常。
                log.error("JWTResponseBodyAdvice error: ", exception);
            }
        }

        return object;
    }
}
